from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.llms import Ollama
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

# 1. 初始化模型
embedding_model = OllamaEmbeddings(model="nomic-embed-text")
llm_model = Ollama(model="llama3.2:3b")

# 2. 加载和分割文档（这里用文本文件示例）
loader = TextLoader("your_document.txt")  # 替换为你的文档路径
documents = loader.load()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200
)
texts = text_splitter.split_documents(documents)

# 3. 创建向量数据库
vectorstore = Chroma.from_documents(
    documents=texts,
    embedding=embedding_model,
    persist_directory="./chroma_db"
)

# 4. 创建检索器
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# 5. 创建自定义提示模板（这是RAG的关键！）
custom_prompt = PromptTemplate(
    template="""基于以下上下文信息，请回答问题。如果你不知道答案，就说不知道，不要编造信息。

上下文：
{context}

问题：{question}
答案：""",
    input_variables=["context", "question"]
)

# 6. 创建RAG链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm_model,
    chain_type="stuff",
    retriever=retriever,
    chain_type_kwargs={"prompt": custom_prompt},
    return_source_documents=True
)

# 7. 提问测试
question = "你的问题在这里"
result = qa_chain.invoke({"query": question})

print("答案:", result["result"])
print("\n来源文档:")
for doc in result["source_documents"]:
    print(f"- {doc.page_content[:100]}...")